"""来自论文Crater edge-based flexible autonomous navigation for planetary landing 
  第7页公式11
"""

from .base import (
    PoseEstimator,
    center_ellipse,
    ransac_pnp,
    pose_calculate,
)
import numpy as np
import cv2
from scipy.optimize import least_squares


class PCCI(PoseEstimator):
    def __init__(self, K, use_pnp_initial=True, *args, **kwargs):
        super().__init__(K, *args, **kwargs)
        self.use_pnp_initial = use_pnp_initial

    def __name__(self):
        return "pcci"

    def forward(self, C_3d, C_2d, *args,**ignore_kwargs):
        if self.use_pnp_initial:
            x_cnt = []
            X_cnt = []
            for c_3d, c_2d in zip(C_3d, C_2d):
                # 求解相机椭圆中心：
                x_cnt.append(center_ellipse(c_2d))
                X_cnt.append(center_ellipse(c_3d))
            # 求解相机位姿
            X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
            x_cnt = np.array(x_cnt)
            ind = ransac_pnp(self.K, X_cnt, x_cnt)
            if ind is None:
                return False, None, None
            R, T = pose_calculate(self.K, X_cnt[ind], x_cnt[ind])
            initial_pose = np.concatenate((cv2.Rodrigues(R)[0], T)).flatten()
        else:
            initial_pose = np.zeros(6)
            X_cnt = []
            for c_3d in C_3d:
                # 求解相机椭圆中心：
                X_cnt.append(center_ellipse(c_3d))
            X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        C_2d = np.array(C_2d)
        return self.estimator(X_cnt, C_2d, initial_pose)

    def estimator(self, database, measurements, initial_pose):
        """
        K: 相机内参
        database: 三维点库
        measurements: 观测到的二维点
        """
        assert len(database) == len(measurements)
        assert measurements.shape[1] == 3
        assert database.shape[1] == 3

        def residual_function(params, X_cnt, C_2d):
            # 定义残差函数
            R = cv2.Rodrigues(params[:3])[0]
            t = params[3:]
            # 求K_i @ R_G^c @ n_i = r_c
            rc = np.linalg.inv(C_2d @ self.K).transpose(0, 2, 1) @ R[:, 2]
            # 几何误差
            rc = rc / np.linalg.norm(rc, axis=1, keepdims=True)
            # 求中心的射线
            # t = -R @ C (C为相机的世界坐标)
            rc_real = X_cnt @ R.T + t
            rc_real = rc_real / np.linalg.norm(rc_real, axis=1, keepdims=True)
            return (rc - rc_real).flatten()

        result = least_squares(
            residual_function, initial_pose, args=(database, measurements)
        )
        if not result.success:
            return False, None, None
        t = result.x[3:]  # 提取平移部分
        R = cv2.Rodrigues(result.x[:3])[0]
        return True, R, t
